import matplotlib.pyplot as plt
import numpy as np

# 基函数  de Boor-Cox 递推定义
def b_basis(p, u_num, U):
    '''
        B_{i,p}(u) : i in range(len(U)), p in range(p+1), u in u_list
    '''
    B = [None]*(len(U)-1-p)
    if p == 0:
        for i in range(len(U)-1):
            if U[i] == U[i+1]:
                B[i] = np.zeros(u_num)
            else:
                B[i] = np.zeros(u_num)
                for k, u in enumerate(np.linspace(U[0], U[-1], u_num, endpoint=True)):
                    B[i][k] = 1 if (U[i] <= u <= U[i+1]) else 0
                    # 注：U[i]<=u<U[i+1]时曲线取不到最后一点，此处随改为闭区间
    else:
        B_ = b_basis(p-1, u_num, U)
        for i in range(len(U)-1-p):
            B[i] = np.zeros(u_num)
            # 计算支撑区间长度
            length1 = U[i+p] - U[i]
            length2 = U[i+p+1] - U[i+1]
            for k, u in enumerate(np.linspace(U[0], U[-1], u_num, endpoint=True)):
                # 定义基函数中的两个系数
                alpha = 0 if not length1 else B_[i][k]/length1 
                beta = 0 if not length2 else B_[i+1][k]/length2
                B[i][k] = alpha*(u-U[i]) + beta*(U[i+p+1]-u)
    return B

def nurbs(U, W, Pts):
    p_num = len(Pts)  # n+1
    u_num = 201 # u插值
    p = len(U)-p_num-1 # 阶数

    # sum_wB
    B = b_basis(p, u_num, U)
    sum_wB = np.zeros(u_num)
    for i in range(p_num):
        sum_wB += W[i]*B[i]

    # 曲线坐标
    points = np.zeros((u_num,3))
    for i in range(p_num):
        Ri = W[i]*B[i] /sum_wB
        points += Ri.reshape(u_num,1) @ [Pts[i]]
    return points


if __name__ == '__main__':
    # 创建画布
    plt.rcParams["font.sans-serif"]=["SimHei"]  # 设置字体
    plt.rcParams["axes.unicode_minus"]=False    # “-”负号的显示
    fig = plt.figure()                          # 空的画布

    # 绘图
    ax = fig.add_subplot(121, projection='3d')  # 创建3d绘图区域
    ax.tick_params(labelsize=8)                 # 刻度字体
    U = [0, 0, 0, 0, 1, 2, 3, 4, 5, 5, 5, 5]
    W = [1, 5, 5, 5, 5, 5, 5, 1] # 权重
    Pts = np.array([(0,0,0),(1,2,1),(5,2,2),(5,4,3),
                    (8,4,4),(4,5,5),(3,5,6),(3,9,7)])
    points = nurbs(U, W, Pts)
    ax.plot(Pts[:,0],Pts[:,1],Pts[:,2])
    ax.plot(points[:,0],points[:,1],points[:,2])
    ax.set_title('Nurbs 曲线')

    # 圆
    ax2 = fig.add_subplot(122, projection='3d')  # 创建3d绘图区域
    ax2.tick_params(labelsize=8)                 # 刻度字体
    x0,y0,z0 = 0,0,0 
    r = 5.0 
    Pts = np.array([( r+x0,   y0, z0), ( r+x0, r+y0, z0), (   x0,  r+y0, z0),
                    (-r+x0, r+y0, z0), (-r+x0,   y0, z0), (-r+x0, -r+y0, z0),
                    (   x0,-r+y0, z0), ( r+x0,-r+y0, z0), ( r+x0,    y0, z0)])
    W = [1, (2**0.5)/2]*4+[1]
    U = [0, 0, 0, 25, 25, 50, 50, 75, 75, 100, 100, 100]
    points = nurbs(U, W, Pts)
    ax2.plot(Pts[:,0],Pts[:,1],Pts[:,2])
    ax2.plot(points[:,0],points[:,1],points[:,2])
    ax2.set_title('Nurbs 圆')

    # 参考线
    def axs(obj, xx=(-1,1),yx=(-1,1),zx=(-1,1)): 
        color = ['r', 'g', 'b']
        for i,t in  enumerate((xx, yx, zx)):
            xs = [(0,0), (0,0), (0,0)]
            f = True if t[0]<0 and t[1]>0 else False
            if f:
                xs[i] = (t[0], 0)
                obj.plot(xs[0], xs[1], xs[2], linestyle=':', color=color[i])
                xs[i] = (0, t[1])
                obj.plot(xs[0], xs[1], xs[2], linestyle='-', color=color[i])
            else: 
                xs[i] = (t[0], t[1])
                obj.plot(xs[0], xs[1], xs[2], linestyle='-', color=color[i])
    axs(ax, (0,5),(0,5),zx=(0,5))
    ax.grid(b=None) # 隐藏网格

    # 显示窗口
    plt.show()